#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
完整的基于权限的患者管理WebSocket服务器
修复数据库结构问题，支持跨网络实时传输
"""

import asyncio
import websockets
import json
import sqlite3
import logging
import time
import os
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

# 配置日志
logging.basicConfig(level=logging.INFO, 
                   format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class CompletePermissionServer:
    def __init__(self, host='0.0.0.0', port=8081):
        self.host = host
        self.port = port
        self.connections = {}
        self.user_sessions = {}
        self.stats = {
            'total_connections': 0,
            'active_connections': 0,
            'messages_processed': 0,
            'start_time': time.time()
        }
        
        # 使用新的数据库文件
        self.db_path = 'complete_patient_management.db'
        self.init_database()
    
    def init_database(self):
        """初始化完整的数据库"""
        # 如果数据库存在就删除重建
        if os.path.exists(self.db_path):
            os.remove(self.db_path)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建完整的用户表
        cursor.execute('''
            CREATE TABLE users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password TEXT NOT NULL,
                real_name TEXT NOT NULL,
                department TEXT NOT NULL,
                user_type TEXT NOT NULL,
                is_active BOOLEAN DEFAULT 1,
                is_approved BOOLEAN DEFAULT 0,
                approved_by INTEGER DEFAULT 0,
                approved_time INTEGER DEFAULT 0,
                phone TEXT,
                registration_date DATETIME DEFAULT CURRENT_TIMESTAMP,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建完整的患者表
        cursor.execute('''
            CREATE TABLE patients (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                gender TEXT NOT NULL,
                age INTEGER,
                phone TEXT,
                id_card TEXT,
                examination_date DATE NOT NULL,
                admission_date DATE,
                diagnosis TEXT,
                notes TEXT,
                reporter TEXT NOT NULL,
                examination_name TEXT,
                department TEXT NOT NULL,
                reported_by INTEGER NOT NULL,
                reporter_type TEXT NOT NULL,
                status TEXT DEFAULT '待检查',
                priority TEXT DEFAULT '普通',
                is_visible BOOLEAN DEFAULT 1,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (reported_by) REFERENCES users(id)
            )
        ''')
        
        # 插入默认用户
        default_users = [
            ('super_admin', 'super123', '大总管理员', '系统管理', 'SUPER_ADMIN', 1, 1),
            ('admin', 'admin123', '总管理员', '系统管理', 'SUPER_ADMIN', 1, 1),
            ('dept_admin_neike', 'dept123', '内科主任', '内科', 'DEPARTMENT_ADMIN', 1, 1),
            ('dept_admin_waike', 'dept123', '外科主任', '外科', 'DEPARTMENT_ADMIN', 1, 1),
            ('user_neike1', 'user123', '内科医生1', '内科', 'USER', 1, 1),
            ('user_waike1', 'user123', '外科医生1', '外科', 'USER', 1, 1)
        ]
        
        for user_data in default_users:
            cursor.execute('''
                INSERT INTO users 
                (username, password, real_name, department, user_type, is_active, is_approved) 
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', user_data)
        
        # 插入测试患者数据
        test_patients = [
            ('张三', '男', 45, '13800138001', '350123198001011234', '2025-08-09', '2025-08-08', '高血压', '需要定期检查', '内科医生1', '血压检查', '内科', 5, 'USER', '待检查', '普通'),
            ('李四', '女', 38, '13800138002', '350123198501021234', '2025-08-09', '2025-08-08', '糖尿病', '需要控制饮食', '内科医生1', '血糖检查', '内科', 5, 'USER', '检查中', '普通'),
            ('王五', '男', 52, '13800138003', '350123197001031234', '2025-08-09', '2025-08-08', '胆结石', '准备手术', '外科医生1', 'B超检查', '外科', 6, 'USER', '待检查', '紧急'),
        ]
        
        for patient_data in test_patients:
            cursor.execute('''
                INSERT INTO patients 
                (name, gender, age, phone, id_card, examination_date, admission_date, diagnosis, notes, 
                 reporter, examination_name, department, reported_by, reporter_type, status, priority)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', patient_data)
        
        conn.commit()
        conn.close()
        logger.info("完整数据库初始化完成，包含测试数据")
    
    def get_db_connection(self):
        """获取数据库连接"""
        return sqlite3.connect(self.db_path, timeout=10.0)
    
    async def register_connection(self, websocket, path):
        """注册新连接"""
        client_id = f"client_{int(time.time())}_{id(websocket)}"
        client_ip = websocket.remote_address[0] if websocket.remote_address else "unknown"
        
        self.connections[client_id] = {
            'websocket': websocket,
            'connected_at': datetime.now(),
            'last_active': datetime.now(),
            'user_info': None,
            'ip_address': client_ip
        }
        
        self.stats['total_connections'] += 1
        self.stats['active_connections'] = len(self.connections)
        
        logger.info(f"新连接: {client_id} from {client_ip}, 当前连接数: {self.stats['active_connections']}")
        
        try:
            # 发送连接确认
            await websocket.send(json.dumps({
                'type': 'CONNECTION_ACK',
                'data': {
                    'client_id': client_id,
                    'server_time': datetime.now().isoformat(),
                    'capabilities': ['permission_based_access', 'realtime_sync', 'department_isolation'],
                    'server_version': '2.0.0-complete'
                }
            }))
            
            # 处理消息
            await self.handle_messages(websocket, client_id)
            
        except websockets.exceptions.ConnectionClosed:
            logger.info(f"连接关闭: {client_id}")
        except Exception as e:
            logger.error(f"连接处理错误 {client_id}: {e}")
        finally:
            await self.unregister_connection(client_id)
    
    async def unregister_connection(self, client_id):
        """注销连接"""
        if client_id in self.connections:
            del self.connections[client_id]
            
            # 清理用户会话
            for user_id, session in list(self.user_sessions.items()):
                if session.get('client_id') == client_id:
                    del self.user_sessions[user_id]
                    logger.info(f"用户 {session.get('real_name', 'Unknown')} 离线")
                    break
            
            self.stats['active_connections'] = len(self.connections)
            logger.info(f"连接注销: {client_id}, 剩余连接数: {self.stats['active_connections']}")
    
    async def handle_messages(self, websocket, client_id):
        """处理消息"""
        async for message in websocket:
            try:
                data = json.loads(message)
                self.stats['messages_processed'] += 1
                
                # 更新最后活动时间
                if client_id in self.connections:
                    self.connections[client_id]['last_active'] = datetime.now()
                
                await self.process_message(data, websocket, client_id)
                
            except json.JSONDecodeError:
                await websocket.send(json.dumps({
                    'type': 'ERROR',
                    'data': {'message': 'Invalid JSON format'}
                }))
            except Exception as e:
                logger.error(f"消息处理错误: {e}")
                await websocket.send(json.dumps({
                    'type': 'ERROR',
                    'data': {'message': str(e)}
                }))
    
    async def process_message(self, data, websocket, client_id):
        """处理具体消息"""
        message_type = data.get('type')
        
        if message_type == 'AUTH':
            await self.handle_auth(data, websocket, client_id)
        elif message_type == 'ADD_PATIENT':
            await self.handle_add_patient(data, websocket, client_id)
        elif message_type == 'GET_PATIENTS':
            await self.handle_get_patients(data, websocket, client_id)
        elif message_type == 'GET_MY_PATIENTS':
            await self.handle_get_my_patients(data, websocket, client_id)
        elif message_type == 'HEARTBEAT':
            await websocket.send(json.dumps({
                'type': 'HEARTBEAT_ACK',
                'timestamp': time.time()
            }))
        elif message_type == 'GET_STATS':
            await self.handle_get_stats(websocket, client_id)
        else:
            logger.warning(f"未知消息类型: {message_type}")
    
    def get_user_permissions(self, user_id):
        """获取用户权限信息"""
        conn = self.get_db_connection()
        cursor = conn.cursor()
        cursor.execute(
            'SELECT id, username, real_name, department, user_type FROM users WHERE id = ? AND is_active = 1',
            (user_id,)
        )
        user_info = cursor.fetchone()
        conn.close()
        
        if user_info:
            return {
                'id': user_info[0],
                'username': user_info[1],
                'real_name': user_info[2],
                'department': user_info[3],
                'user_type': user_info[4]
            }
        return None
    
    def can_access_patient(self, user_permissions, patient_data):
        """检查用户是否可以访问特定患者信息"""
        if not user_permissions:
            return False
        
        user_type = user_permissions['user_type']
        user_dept = user_permissions['department']
        user_id = user_permissions['id']
        
        # 超级管理员可以看到所有信息
        if user_type == 'SUPER_ADMIN':
            return True
        
        # 科室管理员只能看到本科室的信息
        if user_type == 'DEPARTMENT_ADMIN':
            return patient_data.get('department') == user_dept
        
        # 普通用户只能看到自己上报的信息
        if user_type == 'USER':
            return patient_data.get('reported_by') == user_id
        
        return False
    
    async def handle_auth(self, data, websocket, client_id):
        """处理用户认证"""
        user_data = data.get('data', {})
        username = user_data.get('username')
        password = user_data.get('password')
        
        conn = self.get_db_connection()
        cursor = conn.cursor()
        cursor.execute(
            'SELECT id, username, real_name, department, user_type FROM users WHERE username = ? AND password = ? AND is_active = 1 AND is_approved = 1',
            (username, password)
        )
        user_info = cursor.fetchone()
        conn.close()
        
        if user_info:
            user_id = user_info[0]
            
            # 存储用户会话
            self.user_sessions[user_id] = {
                'client_id': client_id,
                'username': user_info[1],
                'real_name': user_info[2],
                'department': user_info[3],
                'user_type': user_info[4],
                'login_time': datetime.now()
            }
            
            # 更新连接信息
            if client_id in self.connections:
                self.connections[client_id]['user_info'] = {
                    'id': user_id,
                    'username': user_info[1],
                    'real_name': user_info[2],
                    'department': user_info[3],
                    'user_type': user_info[4]
                }
            
            await websocket.send(json.dumps({
                'type': 'AUTH_SUCCESS',
                'data': {
                    'user_id': user_id,
                    'user_info': {
                        'id': user_id,
                        'username': user_info[1],
                        'real_name': user_info[2],
                        'department': user_info[3],
                        'user_type': user_info[4]
                    },
                    'permissions': {
                        'can_add_patients': True,
                        'can_view_all_patients': user_info[4] == 'SUPER_ADMIN',
                        'can_view_department_patients': user_info[4] in ['SUPER_ADMIN', 'DEPARTMENT_ADMIN'],
                        'can_approve_users': user_info[4] in ['SUPER_ADMIN', 'DEPARTMENT_ADMIN']
                    },
                    'timestamp': datetime.now().isoformat()
                }
            }))
            
            logger.info(f"用户认证成功: {user_info[2]} ({user_info[4]}) - {user_info[3]}")
        else:
            await websocket.send(json.dumps({
                'type': 'AUTH_FAILED',
                'data': {'message': 'Invalid username or password'}
            }))
    
    async def handle_add_patient(self, data, websocket, client_id):
        """处理添加患者"""
        if client_id not in self.connections or not self.connections[client_id].get('user_info'):
            await websocket.send(json.dumps({
                'type': 'ERROR',
                'data': {'message': 'Not authenticated'}
            }))
            return
        
        user_info = self.connections[client_id]['user_info']
        patient_data = data.get('data', {})
        
        conn = self.get_db_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT INTO patients (name, gender, age, phone, id_card, examination_date, 
                                    diagnosis, notes, reporter, examination_name, department, 
                                    reported_by, reporter_type, admission_date)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                patient_data.get('name'),
                patient_data.get('gender'),
                patient_data.get('age'),
                patient_data.get('phone'),
                patient_data.get('idCard'),
                patient_data.get('examinationDate'),
                patient_data.get('diagnosis'),
                patient_data.get('notes'),
                user_info['real_name'],
                patient_data.get('examinationName'),
                user_info['department'],
                user_info['id'],
                user_info['user_type'],
                patient_data.get('admissionDate')
            ))
            
            patient_id = cursor.lastrowid
            conn.commit()
            
            # 构建患者信息用于广播
            patient_info = {
                'id': patient_id,
                'name': patient_data.get('name'),
                'gender': patient_data.get('gender'),
                'age': patient_data.get('age'),
                'department': user_info['department'],
                'reporter': user_info['real_name'],
                'reported_by': user_info['id'],
                'examination_date': patient_data.get('examinationDate'),
                'diagnosis': patient_data.get('diagnosis'),
                'status': '待检查',
                'timestamp': datetime.now().isoformat()
            }
            
            # 实时广播给有权限的用户
            await self.broadcast_patient_update('NEW_PATIENT_ADDED', patient_info)
            
            await websocket.send(json.dumps({
                'type': 'ADD_PATIENT_SUCCESS',
                'data': {
                    'patient_id': patient_id,
                    'message': '患者信息添加成功',
                    'patient_info': patient_info
                }
            }))
            
            logger.info(f"新患者添加: {patient_data.get('name')} by {user_info['real_name']} ({user_info['department']})")
            
        except Exception as e:
            logger.error(f"患者添加失败: {e}")
            await websocket.send(json.dumps({
                'type': 'ADD_PATIENT_ERROR',
                'data': {'message': f'患者信息添加失败: {str(e)}'}
            }))
        finally:
            conn.close()
    
    async def handle_get_patients(self, data, websocket, client_id):
        """获取患者列表（基于权限）"""
        if client_id not in self.connections or not self.connections[client_id].get('user_info'):
            await websocket.send(json.dumps({
                'type': 'ERROR',
                'data': {'message': 'Not authenticated'}
            }))
            return
        
        user_info = self.connections[client_id]['user_info']
        
        conn = self.get_db_connection()
        cursor = conn.cursor()
        
        # 根据用户类型构建查询
        if user_info['user_type'] == 'SUPER_ADMIN':
            cursor.execute('''
                SELECT id, name, gender, age, phone, examination_date, diagnosis, 
                       notes, reporter, examination_name, department, status, 
                       reported_by, created_at
                FROM patients WHERE is_visible = 1 ORDER BY created_at DESC LIMIT 100
            ''')
        elif user_info['user_type'] == 'DEPARTMENT_ADMIN':
            cursor.execute('''
                SELECT id, name, gender, age, phone, examination_date, diagnosis, 
                       notes, reporter, examination_name, department, status, 
                       reported_by, created_at
                FROM patients WHERE department = ? AND is_visible = 1 
                ORDER BY created_at DESC LIMIT 100
            ''', (user_info['department'],))
        else:
            cursor.execute('''
                SELECT id, name, gender, age, phone, examination_date, diagnosis, 
                       notes, reporter, examination_name, department, status, 
                       reported_by, created_at
                FROM patients WHERE reported_by = ? AND is_visible = 1 
                ORDER BY created_at DESC LIMIT 100
            ''', (user_info['id'],))
        
        patients = []
        for row in cursor.fetchall():
            patients.append({
                'id': row[0],
                'name': row[1],
                'gender': row[2],
                'age': row[3],
                'phone': row[4],
                'examination_date': row[5],
                'diagnosis': row[6],
                'notes': row[7],
                'reporter': row[8],
                'examination_name': row[9],
                'department': row[10],
                'status': row[11],
                'reported_by': row[12],
                'created_at': row[13]
            })
        
        conn.close()
        
        await websocket.send(json.dumps({
            'type': 'PATIENTS_LIST',
            'data': {
                'patients': patients,
                'count': len(patients),
                'user_permissions': {
                    'user_type': user_info['user_type'],
                    'department': user_info['department'],
                    'can_view_all': user_info['user_type'] == 'SUPER_ADMIN'
                },
                'timestamp': datetime.now().isoformat()
            }
        }))
    
    async def handle_get_my_patients(self, data, websocket, client_id):
        """获取我上报的患者"""
        if client_id not in self.connections or not self.connections[client_id].get('user_info'):
            await websocket.send(json.dumps({
                'type': 'ERROR',
                'data': {'message': 'Not authenticated'}
            }))
            return
        
        user_info = self.connections[client_id]['user_info']
        
        conn = self.get_db_connection()
        cursor = conn.cursor()
        cursor.execute('''
            SELECT id, name, gender, age, phone, examination_date, diagnosis, 
                   notes, reporter, examination_name, department, status, created_at
            FROM patients WHERE reported_by = ? AND is_visible = 1 
            ORDER BY created_at DESC
        ''', (user_info['id'],))
        
        patients = []
        for row in cursor.fetchall():
            patients.append({
                'id': row[0],
                'name': row[1],
                'gender': row[2],
                'age': row[3],
                'phone': row[4],
                'examination_date': row[5],
                'diagnosis': row[6],
                'notes': row[7],
                'reporter': row[8],
                'examination_name': row[9],
                'department': row[10],
                'status': row[11],
                'created_at': row[12]
            })
        
        conn.close()
        
        await websocket.send(json.dumps({
            'type': 'MY_PATIENTS_LIST',
            'data': {
                'patients': patients,
                'count': len(patients),
                'timestamp': datetime.now().isoformat()
            }
        }))
    
    async def handle_get_stats(self, websocket, client_id):
        """获取系统统计"""
        if client_id not in self.connections or not self.connections[client_id].get('user_info'):
            await websocket.send(json.dumps({
                'type': 'ERROR',
                'data': {'message': 'Not authenticated'}
            }))
            return
        
        user_info = self.connections[client_id]['user_info']
        
        conn = self.get_db_connection()
        cursor = conn.cursor()
        
        # 获取统计信息
        cursor.execute('SELECT COUNT(*) FROM users WHERE is_approved = 1')
        total_users = cursor.fetchone()[0]
        
        cursor.execute('SELECT COUNT(*) FROM patients WHERE is_visible = 1')
        total_patients = cursor.fetchone()[0]
        
        if user_info['user_type'] == 'SUPER_ADMIN':
            cursor.execute('SELECT department, COUNT(*) FROM patients WHERE is_visible = 1 GROUP BY department')
            dept_stats = dict(cursor.fetchall())
        else:
            cursor.execute('SELECT COUNT(*) FROM patients WHERE department = ? AND is_visible = 1', (user_info['department'],))
            dept_patients = cursor.fetchone()[0]
            dept_stats = {user_info['department']: dept_patients}
        
        conn.close()
        
        await websocket.send(json.dumps({
            'type': 'SYSTEM_STATS',
            'data': {
                'server_stats': {
                    'active_connections': self.stats['active_connections'],
                    'active_users': len(self.user_sessions),
                    'total_messages': self.stats['messages_processed']
                },
                'database_stats': {
                    'total_users': total_users,
                    'total_patients': total_patients,
                    'department_stats': dept_stats
                },
                'timestamp': datetime.now().isoformat()
            }
        }))
    
    async def broadcast_patient_update(self, message_type, patient_info):
        """广播患者信息更新给有权限的用户"""
        message = {
            'type': message_type,
            'data': patient_info
        }
        
        broadcast_count = 0
        for user_id, session in self.user_sessions.items():
            client_id = session.get('client_id')
            if client_id in self.connections:
                try:
                    user_permissions = self.get_user_permissions(user_id)
                    
                    # 检查用户是否有权限查看此患者信息
                    if self.can_access_patient(user_permissions, patient_info):
                        websocket = self.connections[client_id]['websocket']
                        await websocket.send(json.dumps(message))
                        broadcast_count += 1
                        
                        logger.info(f"患者信息已发送给: {session.get('real_name')} ({session.get('user_type')})")
                        
                except Exception as e:
                    logger.error(f"发送患者信息失败: {e}")
        
        logger.info(f"患者信息已广播给 {broadcast_count} 个用户")
    
    async def start_server(self):
        """启动服务器"""
        logger.info(f"启动完整权限管理WebSocket服务器: ws://{self.host}:{self.port}")
        logger.info("权限控制:")
        logger.info("  - 普通用户: 只能看到自己上报的患者信息")
        logger.info("  - 科室管理员: 只能看到本科室的患者信息")
        logger.info("  - 超级管理员: 可以看到所有患者信息")
        logger.info("支持功能: 实时传输、权限隔离、跨网络访问")
        
        async with websockets.serve(
            self.register_connection,
            self.host,
            self.port,
            max_size=1024*1024,
            max_queue=1000,
            ping_interval=30,
            ping_timeout=10
        ):
            logger.info("服务器启动成功，等待连接...")
            await asyncio.Future()

async def main():
    """主函数"""
    server = CompletePermissionServer()
    
    # 启动统计任务
    async def print_stats():
        while True:
            await asyncio.sleep(60)
            logger.info(f"统计信息 - 连接数: {server.stats['active_connections']}, "
                       f"用户数: {len(server.user_sessions)}, "
                       f"消息数: {server.stats['messages_processed']}")
    
    await asyncio.gather(
        server.start_server(),
        print_stats()
    )

if __name__ == "__main__":
    print("🚀 启动完整的基于权限的患者管理系统...")
    print("📊 支持权限控制和跨网络实时传输")
    print("🌐 监听地址: ws://0.0.0.0:8081")
    print("👥 用户权限:")
    print("   - 普通用户: 只能看到自己上报的患者")
    print("   - 科室管理员: 只能看到本科室患者")
    print("   - 超级管理员: 可以看到所有患者")
    print("🎯 实时传输: 患者信息实时推送给有权限的用户")
    print("⏹️ 按Ctrl+C停止服务器")
    
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n⏹️ 服务器已停止")
    except Exception as e:
        print(f"\n❌ 服务器错误: {e}") 